import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']


def plot_results(modelname, stockname, df_index, df_all, generate_data_train, n, generate_data_test, train_end):
    plt.figure(figsize=(18, 6))
    plt.plot(df_index[n:train_end], generate_data_train, label='训练数据价格')
    plt.plot(df_index[train_end:], generate_data_test, c='r', label='预测数据价格')
    plt.plot(df_index[train_end:], df_all[train_end:], c='b', label='真实数据价格')
    plt.title('{modelname}预测{stockname}价格趋势'.format(modelname=modelname,stockname=stockname), fontsize='xx-large', fontweight='heavy')
    plt.xlabel('年份/年', fontdict={'fontsize': 12})
    plt.ylabel('价格/元', fontdict={'fontsize': 12})
    plt.legend()
    plt.show()


def plot_results_custom(modelname, stockname, df_index, df_all, generate_data_test, custom, n):
    plt.rcParams['axes.unicode_minus'] = False
    plt.clf()
    plt.figure(figsize=(18, 6))
    plt.plot(df_index[-250:-1], df_all[-250:-1], label='实际价格')
    plt.plot(df_index[-250:-1], generate_data_test[-250:-1], label='预测价格')
    plt.legend(fontsize=14)
    plt.title('基于{modelname}预测{stockname}价格趋势'.format(modelname=modelname,stockname=stockname), fontsize='xx-large', fontweight='heavy')
    plt.xlabel('年份/年', fontdict={'fontsize': 12})
    plt.ylabel('价格/元', fontdict={'fontsize': 12})
    plt.show()

